#include "eigen/average_pooling_test.h"
#include <Eigen/Dense>
#include "eigen/average_pooling.h"
#include "log.h"
#include "eigen/eigen_equal.h"
#include "tools/range.h"

namespace  ldl_eigen
{
const Eigen::MatrixXf input{
{1,  2,  3,  4,  5, 6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ,20, 21, 22, 23, 24, 25},
{1,  2,  3,  4,  5, 6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ,20, 21, 22, 23, 24, 25},
};

const Eigen::MatrixXf expect_output{
{1.77778, 3.66667, 3.11111, 7.66667, 13, 9.66667, 8.44444, 13.6667, 9.77778},
{1.77778, 3.66667, 3.11111, 7.66667, 13, 9.66667, 8.44444, 13.6667, 9.77778}};

const Eigen::MatrixXf output_gradient{
{9, 8, 7, 6, 5, 4, 3, 2, 1},
{9, 8, 7, 6, 5, 4, 3, 2, 1},
};

const Eigen::MatrixXf expect_input_gradient{
{1, 1.88889, 0.888889, 1.66667, 0.777778, 1.66667, 3.11111, 1.44444, 2.66667, 1.22222, 0.666667, 1.22222, 0.555556, 1, 0.444444, 1, 1.77778, 0.777778, 1.33333, 0.555556, 0.333333, 0.555556, 0.222222, 0.333333, 0.111111},
{1, 1.88889, 0.888889, 1.66667, 0.777778, 1.66667, 3.11111, 1.44444, 2.66667, 1.22222, 0.666667, 1.22222, 0.555556, 1, 0.444444, 1, 1.77778, 0.777778, 1.33333, 0.555556, 0.333333, 0.555556, 0.222222, 0.333333, 0.111111}};

void AveragePoolingTest::test()
{
    const int64_t pooling_size = 3;
    const int64_t stride = 2;
    const int64_t img_height = 5;
    const int64_t img_width = 5;
    const int64_t padding = 1;

    AveragePooling average_pooling(img_width, img_height, pooling_size, stride, padding);

    Eigen::MatrixXf output;
    average_pooling.forward(input, output);
    auto input_gradient =  average_pooling.backward(output_gradient);

    LogInfo() << "input: " << input;
    LogInfo() << "img_height: " << img_height;
    LogInfo() << "img_width: " << img_width;
    LogInfo() << "pooling_size: " << pooling_size;
    LogInfo() << "stride: " << stride;
    LogInfo() << "padding: " << padding;
    LogInfo() << "output_gradient: " << output_gradient;
    // LogInfo() << "output: " << output;
    // LogInfo() << "average_pooling.input_gradient(): " << average_pooling.input_gradient();
    // LogInfo() << "average_pooling.input_gradient().shape: (" << average_pooling.input_gradient().rows() << "," << average_pooling.input_gradient().cols() << ")";

    assert(EigenEqual::equal(output, expect_output, 1e-5));
    assert(EigenEqual::equal(expect_input_gradient, average_pooling.input_gradient()));


    LogInfo() << "AveragePoolingTest test success.";
}
}